package com.example.langchaindemo.base;

import com.example.langchaindemo.Const.APIConst;
import com.example.langchaindemo.config.Langchain4jProperties;
import com.example.langchaindemo.service.AiAssistant;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@RequestMapping("/api")
@RestController
@RequiredArgsConstructor
public class ChatWithRAG {

    final ChatLanguageModel model;

    final AiAssistant assistant;

    final Langchain4jProperties properties;

    @GetMapping("/lowlevel/rag")
    public String lowChat(@RequestParam("message") String message) {
        SystemMessage sys_message = SystemMessage.from("假如你是周杰伦，请用周杰伦的口吻和我对话");
        return model.chat(sys_message, UserMessage.from(message)).aiMessage().text();
    }

    @GetMapping("/highlevel/rag")
    public String highChat(@RequestParam("message") String message) {
        String answer = assistant.chat(message, "说书先生");
        return answer;
    }

    final EmbeddingStore<TextSegment> embeddingStore;
    final EmbeddingModel embeddingModel;

    final EmbeddingStoreContentRetriever retriever;

    @GetMapping("/load")
    public String load(){
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(300, 50);
        Document document = FileSystemDocumentLoader.loadDocument("F:\\GitRepository\\AIModels\\langchain-demo\\src\\main\\resources\\docs\\都市怪谈.txt", new TextDocumentParser());

        List<TextSegment> segments = splitter.split(document);
        /** 批量插入 */
        Response<List<Embedding>> ebeddings = embeddingModel.embedAll(segments);
        embeddingStore.addAll(ebeddings.content(), segments);
        /** 工具插入 */
//        EmbeddingStoreIngestor.builder().embeddingStore(embeddingStore).embeddingModel(embeddingModel).documentSplitter(splitter).build().ingest(document);

        Query qu = Query.from("谁是天明");
        List<Content> retrieve = retriever.retrieve(qu);
        for (Content content : retrieve){
            System.out.println(content);
        }
        return "ok";
    }
}
